package com.easpeed.elevator.utils;

import android.content.Context;
import android.util.Log;
import com.bigyu.utils.Logger;

import com.easpeed.elevator.R;

import ai.onnxruntime.*;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.FloatBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ForkJoinPool;

public class tss {
    private OrtEnvironment env;
    private OrtSession session;
    private FloatBuffer inputBuffer;

    public tss(Context context) throws OrtException, IOException {
        env = OrtEnvironment.getEnvironment();
//        session = env.createSession(loadModel(context));
        session = env.createSession(loadModel(context), new OrtSession.SessionOptions() {{
            setExecutionMode(OrtSession.SessionOptions.ExecutionMode.PARALLEL);
//            addCUDAExecutionProvider(); // 使用 CUDA
        }});

        // 预分配
        this.inputBuffer = FloatBuffer.allocate(250 * 20000);
    }

    private byte[] loadModel(Context context) throws IOException {
        try (InputStream is = context.getResources().openRawResource(R.raw.dvs_model);
             ByteArrayOutputStream buffer = new ByteArrayOutputStream()) {
            byte[] data = new byte[1024];
            int bytesRead;
            while ((bytesRead = is.read(data)) != -1) {
                buffer.write(data, 0, bytesRead);
            }
            return buffer.toByteArray();
        }
    }
    public float[] applyFilter(float[][] inputData) throws OrtException {
        inputBuffer.clear();
        long startTime, endTime, duration;

        // 记录输入数据准备时间
        startTime = System.currentTimeMillis();
        int rows = inputData.length;
        int cols = inputData[0].length;

        // 使用 ForkJoinPool 进行并行填充
        ForkJoinPool.commonPool().submit(() -> {
            for (int i = 0; i < rows; i++) {
                for (int j = 0; j < cols; j++) {
                    inputBuffer.put(inputData[i][j]);
                }
            }
        }).join();
        endTime = System.currentTimeMillis();
        duration = endTime - startTime;
        Logger.d("FirFilters", "Input data preparation time: " + duration + " ms");

        inputBuffer.flip(); // 切换为读取模式

        // 记录张量创建时间
        startTime = System.currentTimeMillis();
        long[] inputNodeDims = new long[] { rows, cols };
        OnnxTensor ont = OnnxTensor.createTensor(env, inputBuffer, inputNodeDims);
        Map<String, OnnxTensor> inputs = new HashMap<>();
        inputs.put("input", ont);
        endTime = System.currentTimeMillis();
        duration = endTime - startTime;
        Logger.d("FirFilters", "Tensor creation time: " + duration + " ms");

        // 记录推理时间
        startTime = System.currentTimeMillis();
        OrtSession.Result result = session.run(inputs);
        float[] outputData = (float[]) result.get(0).getValue();
        endTime = System.currentTimeMillis();
        duration = endTime - startTime;
        Logger.d("FirFilters", "Inference time: " + duration + " ms");

        return outputData;
    }

    public void close() throws OrtException {
        session.close();
        env.close();
    }
}